import asyncio
from typing import List, Union, Any
from loguru import logger
from concurrent.futures import ThreadPoolExecutor

from storage import Neo4jStorage
from llm import LLM
from utils import init_event_loop

from search.prompt import (
    prompt_extract_entities,
    prompt_cypher,
    prompt_categorize,
)
from search.entity import (
    Entity,
    CypherQL,
    TaskCategory
)

from utils import (
    get_today
)

class Task:
    def __init__(self, storage: Neo4jStorage, llm: LLM, question: str):
        self.storage = storage
        self.llm = llm
        self.func_extract_entities = llm.create("Extract Entities", prompt_extract_entities, Entity)
        self.question = question
    
    def _extract_entities(self, question: str) -> Entity:
        return self.func_extract_entities.invoke(input=question, date_time=get_today())

    def retrieve(self) -> str:
        raise NotImplementedError

class SimilarEntity(Task):
    def retrieve(self) -> str:
        entities = self._extract_entities(self.question)

        logger.debug(f"Input {self.question}, get entities {entities}")
        def invoke(entity):
            results = self.storage.get_similiar_entity(entity)
            return "\n".join([f"{r['id']}介绍：{r['summary']}" for r in results])
                       
        futures = [
            invoke(entity)
            for entity in entities.entities
        ]

        results = futures # asyncio.run(asyncio.gather(*futures))

        return "\n".join([r for r in results])

class SummarizeEntity(Task):
    def __init__(self, storage: Neo4jStorage, llm: LLM, question: str):
        super().__init__(storage, llm, question)
        
        self.func_map = self.llm.create(
            "News(Map)", 
            (   
                """当前日期为{date_time}, 根据如下新闻报道，用简单的语言总结有关"{query}"的信息，不要忽略核心的时间、人物、角色、地点、事件\n，严谨作答，如果与问题相关性不大，如果不确定则不输出相关内容，不要编造或者猜测可能性的答案""", 
                ["date_time", "query"]
            )
        )

        self.func_reduce = self.llm.create(
            "News(Reduce)", 
            (
                """当前日期为{date_time}, 请用详实的语言列举出如下问题的回复，必须严谨作答，不要编造或者猜测可能性的答案，如果不确定则不输出相关内容，不超过1000字：{query}""", 
                ["date_time", "query"]
            )
        )

    def retrieve(self) -> str:
        entities = self._extract_entities(self.question)
        logger.debug(f"Input {self.question}, get entities {entities}")
        with ThreadPoolExecutor(max_workers=3) as executor:
            futures = [
                executor.submit(self.storage.get_similiar_entity, entity) 
            for entity in entities.entities]

            similar_entities = [f.result() for f in futures]

        keywords = [[e["id"] for e in entities] for entities in similar_entities]
        # Filter the empty keyword(s)
        keywords = list(filter(lambda x: x is not None and len(x) > 0, keywords))

        resp = self.storage.vector_search(
            self.question, keywords, 
            entities.reported_date_start, 
            entities.reported_date_end, 
            entities.occurred_date_start, 
            entities.occurred_date_end,
            50 # maximium output records
        )

        return self._reduce(self._map(resp))
        
    def _map(self, entries: List[dict]) -> List[str]:
        with ThreadPoolExecutor(max_workers=25) as executor:
            futures = [
                executor.submit(self.func_map.invoke,
                    query=self.question, 
                    date_time=get_today(), 
                    input=f"报道日期：{entry['reported_on']} \n报道正文：{entry['content']}"
                ) 
            for entry in entries]
            results = [f.result() for f in futures]
        # results = loop.run_until_complete(asyncio.gather(*futures))

        return [
            f"报道日期:{entry['reported_on']} - 标题：{entry['title']} - 简介：{res.content}" 
            for res, entry in zip(results, entries)
        ]
    
    def _reduce(self, entries: List[str]) -> str:
        resp = self.func_reduce.invoke(
            query=self.question, 
            date_time=get_today(), 
            input='\n'.join(entries)
        )
        return resp.content

class CypherQuery(Task):
    def __init__(self, storage: Neo4jStorage, llm: LLM, question: str):
        super().__init__(storage, llm, question)
        self.func = llm.create("Generate Cypher", prompt_cypher, CypherQL)

    def retrieve(self) -> str:
        try:
            cql = self.func.invoke(input=self.question, date_time=get_today())
            results = self.storage._query(cql.cypher)
            logger.debug(f"Cypher Generated for {cql} - CQL: {results}")
            return [r for r in results]
        except Exception as e:
            logger.error(f"error in generated cql {cql} - {e} without trying ..")
            return "Cypher 查询失败"

class ChronExecutor:
    def __init__(self):
        self.llm = LLM()
        self.func_category = self.llm.create("Question Categorize", prompt_categorize, TaskCategory)

    def analyze(self, question: str) -> Task:
        storage = Neo4jStorage()

        task = self.func_category.invoke(input=question, date_time=get_today(), context="")

        if task.category == 1:
            return CypherQuery(storage, self.llm, question)
        elif task.category == 2:
            return SimilarEntity(storage, self.llm, question)
        else:
            return SummarizeEntity(storage, self.llm, question)
    
    def retrieve(self, question: str) -> str:
        task = self.analyze(question)
        answer = task.retrieve()

        logger.debug(f"Q:{question} - A: {answer}")

        return f"{answer}"
